/******************************************************************************
 * Copyright 2017 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

#include "modules/drivers/orin_camera/orin_camera_component.h"
#include <sys/epoll.h>
#include <sys/ioctl.h>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <sys/types.h>
#include <opencv2/opencv.hpp>

namespace apollo {
namespace drivers {
namespace orin_camera {

bool OrinCameraComponent::Init() {
  camera_config_ = std::make_shared<Config>();
  if (!apollo::cyber::common::GetProtoFromFile(config_file_path_,
                                               camera_config_.get())) {
    return false;
  }

  if(camera_config_->has_raw_channel_name()){
    pub_raw_channel_ = true;
  }
  if(camera_config_->has_channel_name()){
    pub_image_channel_ = true;
  }

  if(!pub_image_channel_ && !pub_raw_channel_){
    AERROR << "No output channel specified: " << camera_config_->frame_id();
    return false;
  }

  if (!InitSharedMemory()) {
    AERROR << "Init shared memory failed: " << camera_config_->frame_id();
    return false;
  }

  if (pub_image_channel_){
  raw_image_.reset(new CameraImage);

  raw_image_->width = camera_config_->width();
  raw_image_->height = camera_config_->height();
  raw_image_->bytes_per_pixel = camera_config_->bytes_per_pixel();

  if (camera_config_->output_type() == YUYV) {
    raw_image_->image_size = raw_image_->width * raw_image_->height * 2;
  } else if (camera_config_->output_type() == RGB) {
    raw_image_->image_size = raw_image_->width * raw_image_->height * 3;
  }
  if (raw_image_->image_size > MAX_IMAGE_SIZE) {
    AERROR << "image size is too big ,must less than " << MAX_IMAGE_SIZE
           << " bytes.";
    return false;
  }
  raw_image_->is_new = 0;
  // free memory in this struct desturctor
  raw_image_->image =
      reinterpret_cast<char*>(calloc(raw_image_->image_size, sizeof(char)));
  if (raw_image_->image == nullptr) {
    AERROR << "system calloc memory error, size:" << raw_image_->image_size;
    return false;
  }
  writer_ = node_->CreateWriter<Image>(camera_config_->channel_name());
  }

  if (pub_raw_channel_){
    raw_image_for_compress_.reset(new CameraImage);
    raw_image_for_compress_->width = camera_config_->width();
    raw_image_for_compress_->height = camera_config_->height();
    raw_image_for_compress_->bytes_per_pixel = camera_config_->bytes_per_pixel();
    
    if (camera_config_->pixel_format() == "yuyv" ||
        camera_config_->pixel_format() == "uyvy" ||
        camera_config_->pixel_format() == "yuvmono10") {
      raw_image_for_compress_->image_size =
        raw_image_for_compress_->width * raw_image_for_compress_->height * 2;
    } else if (camera_config_->pixel_format() == "mjpeg") {
      AINFO << "Disable sensor raw camera output with format mjpeg";
      raw_image_for_compress_->image_size = 0;
    } else if (camera_config_->pixel_format() == "rgb24") {
      raw_image_for_compress_->image_size =
          raw_image_for_compress_->width * raw_image_for_compress_->height * 3;
    } else {
      AERROR << "Wrong pixel fromat:" << camera_config_->pixel_format()
            << ",must be yuyv | uyvy | mjpeg | yuvmono10 | rgb24";
      return false;
    }

    if (raw_image_for_compress_->image_size == 0) {
      raw_image_for_compress_->image = nullptr;
    } else {
      raw_image_for_compress_->image = reinterpret_cast<char*>(
        calloc(raw_image_for_compress_->image_size, sizeof(char)));
    }
    raw_writer_ = node_->CreateWriter<Image>(camera_config_->raw_channel_name());
  }

  for (int i = 0; i < buffer_size_; ++i) {
    if (pub_image_channel_){
    auto pb_image = std::make_shared<Image>();
    pb_image->mutable_header()->set_frame_id(camera_config_->frame_id());
    pb_image->set_width(raw_image_->width);
    pb_image->set_height(raw_image_->height);
    pb_image->mutable_data()->reserve(raw_image_->image_size);

    if (camera_config_->output_type() == YUYV) {
      pb_image->set_encoding("yuyv");
      pb_image->set_step(2 * raw_image_->width);
    } else if (camera_config_->output_type() == RGB) {
      pb_image->set_encoding("rgb8");
      pb_image->set_step(3 * raw_image_->width);
    }

    pb_image_buffer_.push_back(pb_image);
    }
    if (pub_raw_channel_){
      auto raw_image = std::make_shared<Image>();
      raw_image->mutable_header()->set_frame_id(camera_config_->frame_id());
      raw_image->set_width(raw_image_for_compress_->width);
      raw_image->set_height(raw_image_for_compress_->height);
      raw_image->mutable_data()->reserve(raw_image_for_compress_->image_size);
      raw_image->set_encoding(camera_config_->pixel_format());
      raw_image->set_step(2 * raw_image_for_compress_->width);

      raw_image_buffer_.push_back(raw_image);
    }
  }

  async_result_ = cyber::Async(&OrinCameraComponent::run, this);
  return true;
}

void OrinCameraComponent::run() {
  int ret = 0;
  int port = camera_config_->port();
  int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
  if (listen_fd < 0) { 
    AERROR << "socket faild " << camera_config_->port();
    return;
  }
  int enable = 1;
  ret = setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(enable));
  if (ret < 0) {
    AERROR << "setsockopt faild " << camera_config_->port();
    return;
  }
  struct sockaddr_in server_addr;
  memset(&server_addr, 0, sizeof(server_addr));
  server_addr.sin_family = AF_INET;
  server_addr.sin_addr.s_addr = INADDR_ANY;
  server_addr.sin_port = htons(port);
  ret = bind(listen_fd, (struct sockaddr *)&server_addr, sizeof(server_addr));
  if (ret < 0) {
    AERROR << "bind faild " << camera_config_->port();
    return;
  }
  ret = listen(listen_fd, 3);
  if (ret < 0) {
    AERROR << "listen faild " << camera_config_->port();
    return;
  }
  AINFO << "Server " << camera_config_->frame_id() << " listen on " << camera_config_->port();
  uint64_t trigger_time, sdk_time;
  uint64_t input_image_size = camera_config_->width() * camera_config_->height() * 2;

  running_.exchange(true);
  while (true) {
    // Wait neolix_nv_camera connection
    struct sockaddr_in client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int client_fd = accept(listen_fd, (struct sockaddr *)&client_addr, &client_addr_len);
    if(client_fd < 0) { AERROR << "accept failed " << camera_config_->port(); return;}
    char ip_buf[20];
    memset(ip_buf, 0, sizeof(ip_buf));
    if (inet_ntop(AF_INET, &client_addr.sin_addr, ip_buf, sizeof(ip_buf)) != nullptr) {
      AINFO << "Server" << camera_config_->frame_id() 
      <<" connected by " << ip_buf << ":" << ntohs(client_addr.sin_port);
    }
  
    // For the latest timestamp
    int poll_fd = epoll_create(1);
    if(poll_fd < 0) { AERROR << "epoll_create failed " << camera_config_->port(); return;}
    struct epoll_event reg_event;
    reg_event.events = EPOLLIN;
    reg_event.data.fd = client_fd;
    ret = epoll_ctl(poll_fd, EPOLL_CTL_ADD, client_fd, &reg_event);
    if(ret != 0) { AERROR << "epoll_create failed " << camera_config_->port(); return;}

    std::vector<char> recv_buf;
    // uint64_t read_time;
    while (true) {
      struct epoll_event poll_event;
      ret = epoll_wait(poll_fd, &poll_event, 1, -1);
      if (ret <= 0) {
        AERROR << "Epoll wait failed: " << errno;
        break;
      }
      if(poll_event.data.fd != client_fd) {
        AERROR << "epoll wait fd error " << camera_config_->port(); return;
      }

      // Get the latest timestamp
      int need_recv_len = 0;
      ret = ioctl(client_fd, FIONREAD, &need_recv_len);
      if(ret != 0) { AERROR << "ioctl FIONREAD failed " << camera_config_->port(); return;}
      static const int kRecieveLen = 16;
      if (need_recv_len % kRecieveLen != 0) {
        AINFO << "Receive cache length: " << need_recv_len;
        need_recv_len =
            ((need_recv_len + kRecieveLen - 1) / kRecieveLen) * kRecieveLen;
      }
      uint32_t recv_len = 0;
      recv_buf.resize(need_recv_len);
      while (recv_len < recv_buf.size()) {
        ret = read(client_fd, recv_buf.data() + recv_len,
                   recv_buf.size() - recv_len);
        // read_time = cyber::Time::Now().ToNanosecond();
        if (ret <= 0) {
          AERROR << "Read failed: " << errno;
          break;
        }
        recv_len += ret;
      }
      if (recv_len != recv_buf.size()) {
        AERROR << "Read len failed: " << recv_len << " != " << recv_buf.size();
        break;
      } else if (recv_len == 0) {
        AINFO << "Peer close connection. " << camera_config_->frame_id();
        break;
      }

      // Only process the last timestamp
      for (uint32_t i = 0; i < recv_buf.size(); i += kRecieveLen) {
        trigger_time = *(uint64_t *)(recv_buf.data() + i);
        sdk_time = *(uint64_t *)(recv_buf.data() + i + 8);
        if (i < recv_buf.size() - kRecieveLen) {
          AWARN << "Skip image:" << camera_config_->frame_id() << 
          " trigger:" << trigger_time << " sdk:" << sdk_time;
        }
      }
      AINFO_EVERY(100)  << "Receive image:" << camera_config_->frame_id() 
        << " trigger:" << trigger_time << " sdk:" << sdk_time 
        << " image size:" << input_image_size << " buf size:" << recv_buf.size();
      auto header_time = cyber::Time::Now().ToSecond();

      if (pub_image_channel_){
      auto pb_image = pb_image_buffer_.at(index_);
      pb_image->mutable_header()->set_timestamp_sec(header_time);
      pb_image->set_measurement_time(double(sdk_time) / 1e3);
      pb_image->set_data(raw_image_->image, raw_image_->image_size);

      unsigned char* color_image = new unsigned char[input_image_size];
      memcpy(color_image, image_host_, input_image_size);

      cv::Mat mat_image(static_cast<int>(raw_image_->height), 
                        static_cast<int>(raw_image_->width), CV_8UC2,
                        const_cast<unsigned char*>(color_image));
        cv::Mat tmp_mat;
        cv::cvtColor(mat_image, tmp_mat, cv::COLOR_YUV2RGB_YUYV);
        std::vector<uint8_t> binary_data;
        int total_pixels = raw_image_->height * raw_image_->width * 3;
        binary_data.reserve(total_pixels);

        for (int i = 0; i < raw_image_->height; ++i) {
            for (int j = 0; j < raw_image_->width; ++j) {
                for (int k = 0; k < 3; ++k) {
                    binary_data.push_back(tmp_mat.at<cv::Vec3b>(i, j)[k]);
                }
            }
        }
      pb_image->set_data(binary_data.data(), binary_data.size());
      delete color_image;
      writer_->Write(pb_image);
      }
      if (pub_raw_channel_){
      auto raw_image_for_compress = raw_image_buffer_.at(index_);
      raw_image_for_compress->mutable_header()->set_timestamp_sec(
          header_time);
      raw_image_for_compress->set_measurement_time(double(sdk_time) / 1e3);
      raw_image_for_compress->set_data(raw_image_for_compress_->image,
                                        raw_image_for_compress_->image_size);
      memcpy((void *)raw_image_for_compress->mutable_data()->data(),
               image_host_, input_image_size);
      raw_writer_->Write(raw_image_for_compress);
      }
      index_ = (index_ + 1) % buffer_size_;
    }
    AINFO << "Close connection: " << ip_buf << ":" << ntohs(client_addr.sin_port);
    if (client_fd >= 0) {
      close(client_fd);
    }
    if (poll_fd >= 0) {
      close(poll_fd);
    }
  }
}

OrinCameraComponent::~OrinCameraComponent() {
  if (running_.load()) {
    running_.exchange(false);
    free(raw_image_->image);
    async_result_.wait();
  }
}

bool OrinCameraComponent::InitSharedMemory()  {
  int image_len = 1920 * 1536 * 2;
  key_t shm_key = camera_config_->port();
  shm_id_ = shmget(shm_key, image_len, 0666 | IPC_CREAT);
  if (shm_id_ < 0) {
    AERROR << "shmget failed : " <<  camera_config_->port();
    return false;
  }
  image_host_ = (uint8_t *)shmat(shm_id_, nullptr, 0);
  if (image_host_ == nullptr)  {
    AERROR << "shmat failed : " <<  camera_config_->port();
    return false;
  }

  return true;
}

void OrinCameraComponent::SaveImage(const unsigned char* color_image, std::size_t width,
                            std::size_t height, double timestamp) {
  cv::Mat image(static_cast<int>(height), static_cast<int>(width), CV_8UC2,
                const_cast<unsigned char*>(color_image));
  cv::Mat out(static_cast<int>(height), static_cast<int>(width), CV_8UC3);
  cv::cvtColor(image, out, cv::COLOR_YUV2BGR_YUYV);

  static char path[500];
  snprintf(path, sizeof(path), "/apollo/img/%.6f.png", timestamp);
  // snprintf(path, sizeof(path), "/apollo/img/%.6f.bmp", timestamp);
  cv::imwrite(path, out);
}

// // RGB format
// void OrinCameraComponent::toImage(const unsigned char* color_image, std::size_t width,
//                             std::size_t height, double timestamp) {

//   auto image = std::make_shared<Image>();
//   image->mutable_header()->CopyFrom(image->header());
//   image->set_frame_id(image->frame_id());
//   image->set_measurement_time(image->measurement_time());
//   image->set_height(height);
//   image->set_width(width);
//   image->set_format("rgb8");

//   try {
//   cv::Mat image(static_cast<int>(height), static_cast<int>(width), CV_8UC2,
//                 const_cast<unsigned char*>(color_image));
//     cv::Mat tmp_mat;
//     cv::cvtColor(mat_image, tmp_mat, cv::COLOR_YUV2RGB_YUYV);
//     std::vector<uint8_t> binary_data;
//     int total_pixels = height * width * 3;
//     binary_data.reserve(total_pixels);

//     for (int i = 0; i < height; ++i) {
//         for (int j = 0; j < width; ++j) {
//             for (int k = 0; k < 3; ++k) {
//                 binary_data.push_back(tmp_mat.at<cv::Vec3b>(i, j)[k]);
//             }
//         }
//     }

//     image->set_data(binary_data.data(), binary_data.size());
//   } catch (std::exception& e) {
//     AERROR << "cv::imencode (jpeg) exception :" << e.what();
//     return false;
//   }

//   image_writer_->Write(image);
//   return true;
// }
}  // namespace orin_camera
}  // namespace drivers
}  // namespace apollo
